package com.kanq.handler;

import com.baomidou.mybatisplus.core.toolkit.IdWorker;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.kanq.common.utils.HttpClient;
import com.kanq.common.utils.SseUtil;
import com.kanq.pojo.context.ChatContext;
import com.kanq.pojo.entity.KanqAiFile;
import com.kanq.pojo.entity.KanqAiFileChunk;
import com.kanq.pojo.entity.KanqAiHistory;
import com.kanq.pojo.request.ChatRequest;
import com.kanq.pojo.rerank.RerankRequest;
import com.kanq.pojo.rerank.RerankResponse;
import com.kanq.pojo.response.ChatResponse;
import com.kanq.pojo.response.RefFile;
import com.kanq.service.KanqAiHistoryAnswerService;
import com.kanq.service.KanqAiHistoryQuestionService;
import com.kanq.service.KanqAiHistoryService;
import org.springframework.ai.chat.client.ChatClient;
import reactor.core.Disposable;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;


public class ResponseGeneratorHandler implements ChatHandler {
    private static final ConcurrentHashMap<String, Disposable> taskDisposableMap = new ConcurrentHashMap<>();

    private final ChatClient defaultClient;
    private final KanqAiHistoryService kanqAiHistoryService;
    private final KanqAiHistoryQuestionService kanqAiHistoryQuestionService;
    private final KanqAiHistoryAnswerService kanqAiHistoryAnswerService;
    private final int contextMaxLength;
    private final ChatRequest request;
    private final double scoreThreshold;
    private final String rerankModelName;
    private final String rerankHost;
    private final String rerankKey;

    public ResponseGeneratorHandler(ChatRequest request, ChatClient defaultClient, KanqAiHistoryService kanqAiHistoryService, KanqAiHistoryQuestionService kanqAiHistoryQuestionService, KanqAiHistoryAnswerService kanqAiHistoryAnswerService, int contextMaxLength, double scoreThreshold, String rerankModelName, String rerankHost, String rerankKey) {
        this.request = request;
        this.defaultClient = defaultClient;
        this.kanqAiHistoryService = kanqAiHistoryService;
        this.kanqAiHistoryQuestionService = kanqAiHistoryQuestionService;
        this.kanqAiHistoryAnswerService = kanqAiHistoryAnswerService;
        this.contextMaxLength = contextMaxLength;
        this.scoreThreshold = scoreThreshold;
        this.rerankModelName = rerankModelName;
        this.rerankHost = rerankHost;
        this.rerankKey = rerankKey;
    }

    public static Disposable getDisposableByTaskId(String taskId) {
        return taskDisposableMap.get(taskId);
    }

    @Override
    public void handle(ChatContext context) {
        final boolean[] firstEmptyLine = {true};
        ChatRequest request = context.getRequest();
        List<KanqAiFile> fileWithChunks = context.getFileWithChunks();
        List<KanqAiFileChunk> kanqAiFileChunks = extractChunks(fileWithChunks);

        // 生成模板
        String template = context.getChatDirectlyTemplate().replace("{context}", rerankChunks(request.getQuestion(), kanqAiFileChunks)).replace("{question}", request.getQuestion());

        // 初始化响应
        ChatResponse response = new ChatResponse();
        response.setTaskId(context.getTaskId());
        if (request.getHistoryId() != null) {
            response.setHistoryId(String.valueOf(request.getHistoryId()));
        } else {
            response.setHistoryId(String.valueOf(IdWorker.getId()));
            request.setHistoryId(Long.valueOf(response.getHistoryId()));
        }

        response.setStage("message");
        context.setResponse(response);

        // 流式处理
        StringBuilder sb = new StringBuilder();
        AtomicReference<Long> answerId = new AtomicReference<>();

        // 预先生成问题和标题记录
        CompletableFuture.runAsync(() -> {
            //查看历史内容是否存在
            long count = kanqAiHistoryService.count(Wrappers.<KanqAiHistory>lambdaQuery().eq(KanqAiHistory::getHistoryId, response.getHistoryId()));
            if (count <= 0) {
                String title = kanqAiHistoryService.generateTitle(request.getQuestion());
                kanqAiHistoryService.saveTitle(request.getUserId(), Long.valueOf(response.getHistoryId()), title);
            }
            Long questionId = kanqAiHistoryQuestionService.saveQuestion(request.getQuestion(), Long.valueOf(response.getHistoryId()));
            answerId.set(kanqAiHistoryAnswerService.createAnswerRecord(questionId));
        });

        Disposable disposable = defaultClient.prompt().user(template).advisors(a -> a.param(CONVERSATION_ID, request.getHistoryId())).stream().content().doFinally(signalType -> {
            if (answerId.get() != null) {
                kanqAiHistoryAnswerService.updateAnswer(answerId.get(), sb.toString());
            }
        }).doOnComplete(() -> {
            // 完成处理
            taskDisposableMap.remove(context.getTaskId());
            response.setDone(true);
            response.setStage("message");
            response.setContent("");

            // 设置引用文件
            List<RefFile> refs = new ArrayList<>();
            for (KanqAiFile file : fileWithChunks) {
                refs.add(new RefFile(file.getName(), file.getId()));
            }
            response.setRefs(refs);

            SseUtil.sseSend(context.getEmitter(), response);
        }).subscribe(chunk -> {
            // 检查是否是第一个空行
            if (firstEmptyLine[0] && chunk.equals("\n\n")) {
                firstEmptyLine[0] = true;
            } else {
                firstEmptyLine[0] = false;
                sb.append(chunk);
                response.setContent(chunk);
                SseUtil.sseSend(context.getEmitter(), response);
            }
        }, error -> context.getEmitter().completeWithError(error), () -> context.getEmitter().complete());

        // 存储Disposable
        taskDisposableMap.put(context.getTaskId(), disposable);

        // 错误处理
        context.getEmitter().onError(e -> {
            taskDisposableMap.remove(context.getTaskId());
            if (!disposable.isDisposed()) {
                disposable.dispose();
            }
        });
    }

    String rerankChunks(String query, List<KanqAiFileChunk> chunks) {
        List<String> preResults = new ArrayList<>();
        for (KanqAiFileChunk document : chunks) {
            if (document.getContent().length() > contextMaxLength) {
                continue;
            }
            preResults.add(document.getComment() + document.getContent());
        }

        int totalLength = preResults.stream().mapToInt(String::length).sum();
        if (totalLength <= contextMaxLength) {
            return stringListToString(preResults);
        }

        List<String> results = rerank(query, preResults);
        return stringListToString(results, contextMaxLength);
    }

    public List<String> rerank(String query, List<String> documents) {
        RerankRequest request = new RerankRequest();
        request.setModel(rerankModelName);
        request.setQuery(query);
        request.setTopN(documents.size());
        request.setDocuments(documents);
        request.setDocuments(documents.stream().map(v -> query + v).collect(Collectors.toList()));
        RerankResponse response = HttpClient.sendPostRequest(rerankHost, request, RerankResponse.class, rerankKey);
        if (response == null) {
            return new ArrayList<>();
        }
        var ids = response.getResults(scoreThreshold);
        List<String> results = new ArrayList<>();
        for (var id : ids) {
            results.add(documents.get(id));
        }

        return results;

    }

    String stringListToString(List<String> chunks) {
        StringBuilder sb = new StringBuilder();
        int i = 1;
        for (var c : chunks) {
            sb.append("第").append(i++).append("条信息：").append(c).append("\n\n");
        }

        return sb.toString();
    }

    String stringListToString(List<String> chunks, int maxLength) {
        StringBuilder sb = new StringBuilder();
        int i = 1;
        for (var c : chunks) {
            sb.append("第").append(i).append("条信息：").append(c).append("\n\n");
            if (sb.length() > maxLength) {
                return sb.toString();
            }
            i++;
        }

        return sb.toString();
    }

    List<KanqAiFileChunk> extractChunks(List<KanqAiFile> files) {
        List<KanqAiFileChunk> chunks = new ArrayList<>();
        for (var file : files) {
            chunks.addAll(file.getChunks());
        }
        return chunks;
    }

}